package com.xiashitech.agent.instrumentation.httprequest;

import com.xiashitech.agent.config.StartConfig;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;

public class XSHttpServletRequestWrapper extends HttpServletRequestWrapper {
    /**
     * Constructs a request object wrapping the given request.
     *
     * @param request The request to wrap
     * @throws IllegalArgumentException if the request is null
     */
    public XSHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        this.request = request;
    }

    private byte[] body;
    private HttpServletRequest request;
    @Override
    public ServletInputStream getInputStream() throws IOException {

        if (null == this.body) {
            InputStream inputStream = request.getInputStream();
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            byte[] buffer = new byte[1024];
            int len;
            try {
                while ((len = inputStream.read(buffer)) > -1) {
                    baos.write(buffer, 0, len);
                }
                baos.flush();
            } catch (Throwable e) {
                throw new IOException(e);
            } finally {
                try {
                    inputStream.close();
                } catch (Throwable e) {
                    throw new IOException(e);
                }
            }
            this.body = baos.toByteArray();
        }
        return new XSServletInputStream(new ByteArrayInputStream(body));
    }
    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    /**
     * Get HttpServletRequest Body Multiple Times
     */
    public byte[] getBody() {
        if (null == this.body){
            try {
                getInputStream();
            }catch (IOException e){
                StartConfig.throwableStackPrinter(e);
            }
        }
        return this.body;
    }
}
